from prettytable import PrettyTable
from utils import *
import os
from train import *

# 初始化项目数据结构
train_x, train_y, test_x, test_y = init(
    meta_data_file="./data/meta.json",
    n_path="./data/negative",
    p_path="./data/positive"
) 

if not os.path.exists("model/Haier"):
    os.makedirs("model/Haier")

train_x, train_y = load_abnormal(return_numpy=True, train=True)
test_x, test_y   = load_abnormal(return_numpy=True, train=False)

table = PrettyTable(field_names=["Method name", "cost time(train)", "cost time(test)", "accuracy on test"])

ModelConfig = {
    "SVM"                   : {"parameters" : {"C" : 100},                                                                         "model_class" : ASSVM,         "save_name" : "svc"          },
    "AdaBoost based on SVM" : {"parameters" : {"base_estimator" : SVC(C=100), "n_estimators" : 30, "learning_rate" : 1},           "model_class" : ASAdaBoostSVM, "save_name" : "adaboost_svc" },
    "Bagging based on SVM"  : {"parameters" : {"base_estimator" : SVC(C=100), "n_estimators" : 100, "max_samples" : 0.5},          "model_class" : ASBaggingSVM,  "save_name" : "bagging_svc"  },
    "RandomForest"          : {"parameters" : {"n_estimators" : 200, "criterion" : "gini", "oob_score" : True},                    "model_class" : ASRandomForest,"save_name" : "random_forest"},
    "ExtraTrees"            : {"parameters" : {"n_estimators" : 200, "criterion" : "gini", "bootstrap" : True, "oob_score" : True},"model_class" : ASExtraTrees,  "save_name" : "extra_forest" }
}


for test_name in ModelConfig:
    table.add_row(row_from_train(
        model_name=test_name, model_class=ModelConfig[test_name]["model_class"],parameters=ModelConfig[test_name]["parameters"],
        train_x=train_x, train_y=train_y, test_x=test_x, test_y=test_y, model_path="./model/Haier/{}.joblib".format(ModelConfig[test_name]["save_name"])
    ))


print(Back.GREEN, "DONE", Style.RESET_ALL,  Fore.GREEN + "Finish Initialization, Ready to go :D" + Style.RESET_ALL)
print(table)

"""
    This is a simple demo and something I wanna tell you:
    - All the registered model has a prefix 'AS', which means anomalous detection
    - All the JobLib class share the same method to load, save, train, test
"""

for test_name in ModelConfig:
    print(Back.BLUE, "TEST", Style.RESET_ALL, "testing {}".format(test_name), Style.RESET_ALL)
    model = ModelConfig[test_name]["model_class"]()
    model.load_model("./model/Haier/{}.joblib".format(ModelConfig[test_name]["save_name"]))
    print(Back.GREEN, "Test sample:", Style.RESET_ALL, model.predict_wav_file("./data/negative/1.wav"))
    model.report(test_x, test_y)
    print(Back.BLUE, "acc" , Style.RESET_ALL, Fore.BLUE, model.score(test_x, test_y),   Style.RESET_ALL)
    print(Back.BLUE, "auc ", Style.RESET_ALL, Fore.BLUE, model.roc_auc(test_x, test_y), Style.RESET_ALL)

print(Back.GREEN, "DONE", Style.RESET_ALL, Fore.GREEN, "everything is done :D", Style.RESET_ALL)